import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from solve_poisson_equation_direct import solve_poisson_equation_direct  # Assuming you saved the previous function in a module

# Clear previous figures (equivalent to MATLAB's clear command)
plt.close('all')

# Set the number of points in the x and y directions
Nx = 32
Ny = 32

# Solve the Poisson equation using the direct method
x, y, u = solve_poisson_equation_direct(Nx, Ny)

# Plot the solution
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# Surface plot
surf = ax.plot_surface(x, y, u, cmap='viridis', edgecolor='none')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('u')
ax.set_title('Solution of the Poisson Equation')

# Add a color bar which maps values to colors
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()